'''
Author: tomwoo tom.woo@outlook.com
Date: 2025-07-27 23:26:12
LastEditors: tomwoo tom.woo@outlook.com
LastEditTime: 2025-08-20 00:10:10
FilePath: /Weather_Forecast_Agent/forecast_tools/plot.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

from typing import Optional

import cartopy.crs as ccrs
import matplotlib.animation as animation
import matplotlib.font_manager as fm
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize


font_dirs = ["/usr/share/fonts/opentype"] # The paths to the custom font files.
font_files = fm.findSystemFonts(fontpaths=font_dirs)

for font_file in font_files:
    fm.fontManager.addfont(font_file)

# chinese_fonts = [f.name for f in font_manager.fontManager.ttflist if 'CJK' in f.name]
# print(chinese_fonts) # 输出所有中文字体名称列表

plt.rcParams['font.sans-serif'] = ['Noto Sans CJK JP']


def align_coords(*coords: float, factor: int = 4) -> list[float]:
    return [round(c * factor) / factor for c in coords]

def animate_local(
    da: xr.DataArray,
    start_time: np.datetime64,
    timezone_delta: np.timedelta64,
    loc_lat: float,
    loc_lon: float,
    max_frames: int = 32,
    pcolormesh_kwargs: Optional[dict] = None,
    cb_label: Optional[str] = None,
    interval: int = 500
) -> animation.FuncAnimation:
    print("Creating animation...", end="")
    if pcolormesh_kwargs is None:
        pcolormesh_kwargs = {}

    da = da.sel(time=start_time).drop_vars("time")

    proj = ccrs.PlateCarree(central_longitude=loc_lon)
    fig, ax = plt.subplots(subplot_kw={"projection": proj}, figsize=(6, 5))

    loc_lat, loc_lon = align_coords(loc_lat, loc_lon)
    loc_da = da.sel(
        lon=slice(loc_lon - 12, loc_lon + 12), lat=slice(loc_lat + 12, loc_lat - 12)
    )
    im = ax.pcolormesh(
        loc_da.lon,
        loc_da.lat,
        loc_da.isel(lead_time=0),
        transform=ccrs.PlateCarree(),
        **{"cmap": "coolwarm", **pcolormesh_kwargs},
    )

    ax.set_title(f"{start_time} (lead time: 0 h)")
    ax.coastlines(linewidth=1.0)

    def animate(step: int) -> None:
        print(".", end="")
        im.set_array(loc_da.isel(lead_time=step))
        utc_step_time = np.datetime_as_string(start_time + loc_da.lead_time[step], unit="s")
        local_step_time = np.datetime_as_string(start_time + loc_da.lead_time[step] + timezone_delta, unit="s")
        str_a = f"世界标准时间：{utc_step_time.replace('T', ' ')} (lead time: {6 * step} h)"
        str_b = f"当地日期时间：{local_step_time.replace('T', ' ')} (lead time: {6 * step} h)"
        ax.set_title(str_a + '\n' + str_b)

    # Inset colorbar
    vmin = pcolormesh_kwargs.get("vmin", loc_da.values.min())
    vmax = pcolormesh_kwargs.get("vmax", loc_da.values.max())
    norm = Normalize(vmin, vmax, clip=True)
    cmap = pcolormesh_kwargs.get("cmap", "coolwarm")

    cax = ax.inset_axes([0.04, 0.02, 0.015, 0.4])
    fig.colorbar(ScalarMappable(norm=norm, cmap=cmap), cax=cax, label=cb_label)

    fig.tight_layout()
    plt.close("all")

    return animation.FuncAnimation(
        fig,
        animate,
        min(max_frames, len(loc_da.lead_time)),
        blit=False,
        repeat=True,
        interval=interval,
    )

# end of file
